// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <nebula/io/compressed.h>

#include <algorithm>
#include <cstring>
#include <memory>
#include <mutex>
#include <string>
#include <utility>

#include <nebula/core/buffer.h>
#include <nebula/io/util_internal.h>
#include <nebula/core/memory_pool.h>
#include <turbo/utility/status.h>
#include <nebula/compression/compression.h>
#include <turbo/log/logging.h>

namespace nebula::io {
    // ----------------------------------------------------------------------
    // CompressedOutputStream implementation

    class CompressedOutputStream::Impl {
    public:
        Impl(MemoryPool *pool, const std::shared_ptr<OutputStream> &raw)
            : pool_(pool), raw_(raw), is_open_(false), compressed_pos_(0), total_pos_(0) {
        }

        turbo::Status init(Codec *codec) {
            TURBO_MOVE_OR_RAISE(compressor_, codec->MakeCompressor());
            TURBO_MOVE_OR_RAISE(compressed_, allocate_resizable_buffer(kChunkSize, pool_));
            compressed_pos_ = 0;
            is_open_ = true;
            return turbo::OkStatus();
        }

        turbo::Result<int64_t> tell() const {
            std::lock_guard<std::mutex> guard(lock_);
            return total_pos_;
        }

        std::shared_ptr<OutputStream> raw() const { return raw_; }

        turbo::Status FlushCompressed() {
            if (compressed_pos_ > 0) {
                TURBO_RETURN_NOT_OK(raw_->write(compressed_->data(), compressed_pos_));
                compressed_pos_ = 0;
            }
            return turbo::OkStatus();
        }

        turbo::Status write(const void *data, int64_t nbytes) {
            std::lock_guard<std::mutex> guard(lock_);

            auto input = reinterpret_cast<const uint8_t *>(data);
            while (nbytes > 0) {
                int64_t input_len = nbytes;
                int64_t output_len = compressed_->size() - compressed_pos_;
                uint8_t *output = compressed_->mutable_data() + compressed_pos_;
                TURBO_MOVE_OR_RAISE(auto result,
                                    compressor_->Compress(input_len, input, output_len, output));
                compressed_pos_ += result.bytes_written;

                if (result.bytes_read == 0) {
                    // Not enough output, try to flush it and retry
                    if (compressed_pos_ > 0) {
                        TURBO_RETURN_NOT_OK(FlushCompressed());
                        output_len = compressed_->size() - compressed_pos_;
                        output = compressed_->mutable_data() + compressed_pos_;
                        TURBO_MOVE_OR_RAISE(
                            result, compressor_->Compress(input_len, input, output_len, output));
                        compressed_pos_ += result.bytes_written;
                    }
                }
                input += result.bytes_read;
                nbytes -= result.bytes_read;
                total_pos_ += result.bytes_read;
                if (compressed_pos_ == compressed_->size()) {
                    // Output buffer full, flush it
                    TURBO_RETURN_NOT_OK(FlushCompressed());
                }
                if (result.bytes_read == 0) {
                    // Need to enlarge output buffer
                    TURBO_RETURN_NOT_OK(compressed_->resize(compressed_->size() * 2));
                }
            }
            return turbo::OkStatus();
        }

        turbo::Status flush() {
            std::lock_guard<std::mutex> guard(lock_);

            while (true) {
                // flush compressor
                int64_t output_len = compressed_->size() - compressed_pos_;
                uint8_t *output = compressed_->mutable_data() + compressed_pos_;
                TURBO_MOVE_OR_RAISE(auto result, compressor_->flush(output_len, output));
                compressed_pos_ += result.bytes_written;

                // flush compressed output
                TURBO_RETURN_NOT_OK(FlushCompressed());

                if (result.should_retry) {
                    // Need to enlarge output buffer
                    TURBO_RETURN_NOT_OK(compressed_->resize(compressed_->size() * 2));
                } else {
                    break;
                }
            }
            return turbo::OkStatus();
        }

        turbo::Status FinalizeCompression() {
            while (true) {
                // Try to end compressor
                int64_t output_len = compressed_->size() - compressed_pos_;
                uint8_t *output = compressed_->mutable_data() + compressed_pos_;
                TURBO_MOVE_OR_RAISE(auto result, compressor_->End(output_len, output));
                compressed_pos_ += result.bytes_written;

                // flush compressed output
                TURBO_RETURN_NOT_OK(FlushCompressed());

                if (result.should_retry) {
                    // Need to enlarge output buffer
                    TURBO_RETURN_NOT_OK(compressed_->resize(compressed_->size() * 2));
                } else {
                    // Done
                    break;
                }
            }
            return turbo::OkStatus();
        }

        turbo::Status close() {
            std::lock_guard<std::mutex> guard(lock_);

            if (is_open_) {
                is_open_ = false;
                TURBO_RETURN_NOT_OK(FinalizeCompression());
                return raw_->close();
            } else {
                return turbo::OkStatus();
            }
        }

        turbo::Status abort() {
            std::lock_guard<std::mutex> guard(lock_);

            if (is_open_) {
                is_open_ = false;
                return raw_->abort();
            } else {
                return turbo::OkStatus();
            }
        }

        bool closed() {
            std::lock_guard<std::mutex> guard(lock_);
            return !is_open_;
        }

    private:
        // write 64 KB compressed data at a time
        static const int64_t kChunkSize = 64 * 1024;

        MemoryPool *pool_;
        std::shared_ptr<OutputStream> raw_;
        bool is_open_;
        std::shared_ptr<Compressor> compressor_;
        std::shared_ptr<ResizableBuffer> compressed_;
        int64_t compressed_pos_;
        // Total number of bytes compressed
        int64_t total_pos_;

        mutable std::mutex lock_;
    };

    turbo::Result<std::shared_ptr<CompressedOutputStream> > CompressedOutputStream::create(
        Codec *codec, const std::shared_ptr<OutputStream> &raw, MemoryPool *pool) {
        // CAUTION: codec is not owned
        std::shared_ptr<CompressedOutputStream> res(new CompressedOutputStream);
        res->impl_.reset(new Impl(pool, raw));
        TURBO_RETURN_NOT_OK(res->impl_->init(codec));
        return res;
    }

    CompressedOutputStream::~CompressedOutputStream() { internal::CloseFromDestructor(this); }

    turbo::Status CompressedOutputStream::close() { return impl_->close(); }

    turbo::Status CompressedOutputStream::abort() { return impl_->abort(); }

    bool CompressedOutputStream::closed() const { return impl_->closed(); }

    turbo::Result<int64_t> CompressedOutputStream::tell() const { return impl_->tell(); }

    turbo::Status CompressedOutputStream::write(const void *data, int64_t nbytes) {
        return impl_->write(data, nbytes);
    }

    turbo::Status CompressedOutputStream::flush() { return impl_->flush(); }

    std::shared_ptr<OutputStream> CompressedOutputStream::raw() const { return impl_->raw(); }

    // ----------------------------------------------------------------------
    // CompressedInputStream implementation

    class CompressedInputStream::Impl {
    public:
        Impl(MemoryPool *pool, const std::shared_ptr<InputStream> &raw)
            : pool_(pool),
              raw_(raw),
              is_open_(true),
              supports_zero_copy_from_raw_(raw_->supports_zero_copy()),
              compressed_pos_(0),
              decompressed_pos_(0),
              fresh_decompressor_(false),
              total_pos_(0) {
        }

        turbo::Status init(Codec *codec) {
            TURBO_MOVE_OR_RAISE(decompressor_, codec->MakeDecompressor());
            fresh_decompressor_ = true;
            return turbo::OkStatus();
        }

        turbo::Status close() {
            if (is_open_) {
                is_open_ = false;
                return raw_->close();
            } else {
                return turbo::OkStatus();
            }
        }

        turbo::Status abort() {
            if (is_open_) {
                is_open_ = false;
                return raw_->abort();
            } else {
                return turbo::OkStatus();
            }
        }

        bool closed() const { return !is_open_; }

        turbo::Result<int64_t> tell() const { return total_pos_; }

        // read compressed data if necessary
        turbo::Status EnsureCompressedData() {
            int64_t compressed_avail = compressed_buffer_available();
            if (compressed_avail == 0) {
                // Ensure compressed_ buffer is allocated with kChunkSize.
                if (!supports_zero_copy_from_raw_) {
                    if (compressed_for_non_zero_copy_ == nullptr) {
                        TURBO_MOVE_OR_RAISE(compressed_for_non_zero_copy_,
                                            allocate_resizable_buffer(kChunkSize, pool_));
                    } else if (compressed_for_non_zero_copy_->size() != kChunkSize) {
                        TURBO_RETURN_NOT_OK(
                            compressed_for_non_zero_copy_->resize(kChunkSize, /*shrink_to_fit=*/false));
                    }
                    TURBO_MOVE_OR_RAISE(
                        int64_t read_size,
                        raw_->read(kChunkSize,
                            compressed_for_non_zero_copy_->mutable_data_as<void>()));
                    if (read_size != compressed_for_non_zero_copy_->size()) {
                        TURBO_RETURN_NOT_OK(
                            compressed_for_non_zero_copy_->resize(read_size, /*shrink_to_fit=*/false));
                    }
                    compressed_ = compressed_for_non_zero_copy_;
                } else {
                    TURBO_MOVE_OR_RAISE(compressed_, raw_->read(kChunkSize));
                }
                compressed_pos_ = 0;
            }
            return turbo::OkStatus();
        }

        // Decompress some data from the compressed_ buffer into decompressor_.
        // Call this function only if the decompressed_ buffer is fully consumed.
        turbo::Status DecompressData() {
            // compressed_buffer_available() could be 0 here because there might
            // still be some decompressed data left to emit even though the compressed
            // data was entirely consumed (especially if the expansion factor is large)
            DKCHECK_NE(compressed_->data(), nullptr);
            DKCHECK_EQ(0, decompressed_buffer_available());

            int64_t decompress_size = kDecompressSize;

            while (true) {
                if (decompressed_ == nullptr) {
                    TURBO_MOVE_OR_RAISE(decompressed_,
                                        allocate_resizable_buffer(decompress_size, pool_));
                } else {
                    // Shrinking the buffer if it's already large enough
                    TURBO_RETURN_NOT_OK(decompressed_->resize(decompress_size, /*shrink_to_fit=*/true));
                }
                decompressed_pos_ = 0;

                int64_t input_len = compressed_->size() - compressed_pos_;
                const uint8_t *input = compressed_->data() + compressed_pos_;
                int64_t output_len = decompressed_->size();
                uint8_t *output = decompressed_->mutable_data();

                TURBO_MOVE_OR_RAISE(
                    auto result, decompressor_->Decompress(input_len, input, output_len, output));
                compressed_pos_ += result.bytes_read;
                if (result.bytes_read > 0) {
                    fresh_decompressor_ = false;
                }
                if (result.bytes_written > 0 || !result.need_more_output || input_len == 0) {
                    // Not calling shrink_to_fit here because we're likely to reusing the buffer.
                    TURBO_RETURN_NOT_OK(
                        decompressed_->resize(result.bytes_written, /*shrink_to_fit=*/false));
                    break;
                }
                DKCHECK_EQ(result.bytes_written, 0);
                // Need to enlarge output buffer
                decompress_size *= 2;
            }
            return turbo::OkStatus();
        }

        // Copying a given number of bytes from the decompressed_ buffer.
        int64_t ReadFromDecompressed(int64_t nbytes, uint8_t *out) {
            int64_t readable = decompressed_ ? (decompressed_->size() - decompressed_pos_) : 0;
            int64_t read_bytes = std::min(readable, nbytes);

            if (read_bytes > 0) {
                memcpy(out, decompressed_->data() + decompressed_pos_, read_bytes);
                decompressed_pos_ += read_bytes;
            }

            return read_bytes;
        }

        // Try to feed more data into the decompressed_ buffer.
        // Returns whether there is more data to read.
        turbo::Result<bool> RefillDecompressed() {
            // First try to read data from the decompressor, unless we haven't read any
            // compressed data yet.
            if (compressed_ && compressed_->size() != 0) {
                if (decompressor_->IsFinished()) {
                    // We just went over the end of a previous compressed stream.
                    TURBO_RETURN_NOT_OK(decompressor_->reset());
                    fresh_decompressor_ = true;
                }
                TURBO_RETURN_NOT_OK(DecompressData());
            }
            int64_t decompress_avail = decompressed_buffer_available();
            if (decompress_avail == 0) {
                // Got nothing from existing `compressed_` and `decompressor_`,
                // need to read more compressed data.
                TURBO_RETURN_NOT_OK(EnsureCompressedData());
                if (compressed_buffer_available() == 0) {
                    // No more data to decompress
                    if (!fresh_decompressor_ && !decompressor_->IsFinished()) {
                        return turbo::io_error("Truncated compressed stream");
                    }
                    return false;
                }
                TURBO_RETURN_NOT_OK(DecompressData());
            }
            return true;
        }

        turbo::Result<int64_t> read(int64_t nbytes, void *out) {
            auto *out_data = reinterpret_cast<uint8_t *>(out);

            int64_t total_read = 0;
            bool decompressor_has_data = true;

            while (nbytes - total_read > 0 && decompressor_has_data) {
                total_read += ReadFromDecompressed(nbytes - total_read, out_data + total_read);

                if (nbytes == total_read) {
                    break;
                }

                // At this point, no more decompressed data remains, so we need to
                // decompress more
                TURBO_MOVE_OR_RAISE(decompressor_has_data, RefillDecompressed());
            }

            total_pos_ += total_read;
            return total_read;
        }

        turbo::Result<std::shared_ptr<Buffer> > read(int64_t nbytes) {
            TURBO_MOVE_OR_RAISE(auto buf, allocate_resizable_buffer(nbytes, pool_));
            TURBO_MOVE_OR_RAISE(int64_t bytes_read, read(nbytes, buf->mutable_data()));
            TURBO_RETURN_NOT_OK(buf->resize(bytes_read));
            // R build with openSUSE155 requires an explicit shared_ptr construction
            return std::shared_ptr<Buffer>(std::move(buf));
        }

        const std::shared_ptr<InputStream> &raw() const { return raw_; }

    private:
        int64_t compressed_buffer_available() const {
            return compressed_ ? compressed_->size() - compressed_pos_ : 0;
        }

        int64_t decompressed_buffer_available() const {
            return decompressed_ ? decompressed_->size() - decompressed_pos_ : 0;
        }

    private:
        // read 64 KB compressed data at a time
        static const int64_t kChunkSize = 64 * 1024;
        // Decompress 1 MB at a time
        static const int64_t kDecompressSize = 1024 * 1024;

        MemoryPool *pool_;
        std::shared_ptr<InputStream> raw_;
        bool is_open_;
        const bool supports_zero_copy_from_raw_;
        std::shared_ptr<Decompressor> decompressor_;
        // If `raw_->supports_zero_copy()`, this buffer would not allocate memory.
        // Otherwise, this buffer would allocate `kChunkSize` memory and read data from
        // `raw_`.
        std::shared_ptr<ResizableBuffer> compressed_for_non_zero_copy_;
        std::shared_ptr<Buffer> compressed_;
        // Position in compressed buffer
        int64_t compressed_pos_;
        std::shared_ptr<ResizableBuffer> decompressed_;
        // Position in decompressed buffer
        int64_t decompressed_pos_;
        // True if the decompressor hasn't read any data yet.
        bool fresh_decompressor_;
        // Total number of bytes decompressed
        int64_t total_pos_;
    };

    turbo::Result<std::shared_ptr<CompressedInputStream> > CompressedInputStream::create(
        Codec *codec, const std::shared_ptr<InputStream> &raw, MemoryPool *pool) {
        // CAUTION: codec is not owned
        std::shared_ptr<CompressedInputStream> res(new CompressedInputStream);
        res->impl_.reset(new Impl(pool, raw));
        TURBO_RETURN_NOT_OK(res->impl_->init(codec));
        return res;
    }

    CompressedInputStream::~CompressedInputStream() { internal::CloseFromDestructor(this); }

    turbo::Status CompressedInputStream::DoClose() { return impl_->close(); }

    turbo::Status CompressedInputStream::DoAbort() { return impl_->abort(); }

    bool CompressedInputStream::closed() const { return impl_->closed(); }

    turbo::Result<int64_t> CompressedInputStream::DoTell() const { return impl_->tell(); }

    turbo::Result<int64_t> CompressedInputStream::DoRead(int64_t nbytes, void *out) {
        return impl_->read(nbytes, out);
    }

    turbo::Result<std::shared_ptr<Buffer> > CompressedInputStream::DoRead(int64_t nbytes) {
        return impl_->read(nbytes);
    }

    std::shared_ptr<InputStream> CompressedInputStream::raw() const { return impl_->raw(); }

    turbo::Result<std::shared_ptr<const KeyValueMetadata> > CompressedInputStream::read_metadata() {
        return impl_->raw()->read_metadata();
    }

    Future<std::shared_ptr<const KeyValueMetadata> > CompressedInputStream::read_metadata_async(
        const IOContext &io_context) {
        return impl_->raw()->read_metadata_async(io_context);
    }
} // namespace nebula::io
